#!/bin/bash

# 5GC Testbed Startup Script
# This script automates the startup of all 5GC components in the correct order

set -e  # Exit on any error

# Configuration
VM1="192.168.1.206"  # Core Network
VM2="192.168.1.219"  # UPF
VM3="192.168.1.221"  # gNodeB
VM4="192.168.1.205"  # UE

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Logging function
log() {
    echo -e "${GREEN}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}"
}

warn() {
    echo -e "${YELLOW}[$(date +'%Y-%m-%d %H:%M:%S')] WARNING: $1${NC}"
}

error() {
    echo -e "${RED}[$(date +'%Y-%m-%d %H:%M:%S')] ERROR: $1${NC}"
}

# Function to check if a service is running
check_service() {
    local service=$1
    local vm=$2
    
    if ssh -o ConnectTimeout=5 $vm "systemctl is-active --quiet $service" 2>/dev/null; then
        return 0
    else
        return 1
    fi
}

# Function to check if a process is running
check_process() {
    local process=$1
    local vm=$2
    
    if ssh -o ConnectTimeout=5 $vm "pgrep -f '$process'" >/dev/null 2>&1; then
        return 0
    else
        return 1
    fi
}

# Function to wait for service to be ready
wait_for_service() {
    local service=$1
    local vm=$2
    local max_attempts=30
    local attempt=1
    
    log "Waiting for $service on $vm to be ready..."
    
    while [ $attempt -le $max_attempts ]; do
        if check_service $service $vm; then
            log "✓ $service on $vm is ready"
            return 0
        fi
        
        echo -n "."
        sleep 2
        ((attempt++))
    done
    
    error "$service on $vm failed to start within $((max_attempts * 2)) seconds"
    return 1
}

# Function to wait for process to be ready
wait_for_process() {
    local process=$1
    local vm=$2
    local max_attempts=30
    local attempt=1
    
    log "Waiting for $process on $vm to be ready..."
    
    while [ $attempt -le $max_attempts ]; do
        if check_process $process $vm; then
            log "✓ $process on $vm is ready"
            return 0
        fi
        
        echo -n "."
        sleep 2
        ((attempt++))
    done
    
    error "$process on $vm failed to start within $((max_attempts * 2)) seconds"
    return 1
}

# Pre-flight checks
log "Performing pre-flight checks..."

# Check VM connectivity
for vm in $VM1 $VM2 $VM3 $VM4; do
    if ping -c 1 -W 5 $vm >/dev/null 2>&1; then
        log "✓ VM $vm is reachable"
    else
        error "VM $vm is not reachable"
        exit 1
    fi
done

# Check if services are already running
log "Checking current service status..."

# Stop any existing services to ensure clean startup
log "Stopping any existing services..."

ssh $VM1 "sudo systemctl stop open5gs-* 2>/dev/null || true"
ssh $VM2 "sudo systemctl stop open5gs-upfd 2>/dev/null || true"
ssh $VM3 "pkill -f nr-gnb 2>/dev/null || true"
ssh $VM4 "pkill -f nr-ue 2>/dev/null || true"

sleep 5

# Step 1: Start Core Network Services (VM1)
log "Starting core network services on VM1..."

# Start MongoDB first
log "Starting MongoDB..."
ssh $VM1 "sudo systemctl start mongod"
wait_for_service "mongod" $VM1

# Start core services in order
services=("open5gs-udrd" "open5gs-udmd" "open5gs-pcfd" "open5gs-smfd" "open5gs-amfd")

for service in "${services[@]}"; do
    log "Starting $service..."
    ssh $VM1 "sudo systemctl start $service"
    wait_for_service $service $VM1
    sleep 2
done

# Step 2: Start User Plane Function (VM2)
log "Starting UPF on VM2..."

ssh $VM2 "sudo systemctl start open5gs-upfd"
wait_for_service "open5gs-upfd" $VM2

# Step 3: Start Radio Access Network (VM3)
log "Starting gNodeB on VM3..."

# Navigate to UERANSIM directory and start gNodeB
ssh $VM3 << 'EOF'
cd /path/to/ueransim  # Update this path as needed
nohup ./nr-gnb -c config/gnb.yaml > gnb.log 2>&1 &
EOF

wait_for_process "nr-gnb" $VM3

# Step 4: Start User Equipment (VM4)
log "Starting UE on VM4..."

# Navigate to UERANSIM directory and start UE
ssh $VM4 << 'EOF'
cd /path/to/ueransim  # Update this path as needed
nohup ./nr-ue -c config/ue.yaml > ue.log 2>&1 &
EOF

wait_for_process "nr-ue" $VM4

# Final verification
log "Performing final verification..."

# Check all services
log "Verifying all services are running..."

# Core services
for service in "${services[@]}"; do
    if check_service $service $VM1; then
        log "✓ $service is running"
    else
        error "$service is not running"
    fi
done

# UPF
if check_service "open5gs-upfd" $VM2; then
    log "✓ UPF is running"
else
    error "UPF is not running"
fi

# UERANSIM processes
if check_process "nr-gnb" $VM3; then
    log "✓ gNodeB is running"
else
    error "gNodeB is not running"
fi

if check_process "nr-ue" $VM4; then
    log "✓ UE is running"
else
    error "UE is not running"
fi

# Check network connectivity
log "Verifying network connectivity..."

# Check SCTP connections
if ssh $VM3 "ss -tan | grep '192.168.1.10:38412' | grep ESTAB" >/dev/null 2>&1; then
    log "✓ gNodeB-AMF SCTP association established"
else
    warn "gNodeB-AMF SCTP association not yet established"
fi

# Check listening ports
log "Verifying listening ports..."

# Core network ports
if ssh $VM1 "ss -tlnp | grep -E '(7777|38412|8805|2123)'" >/dev/null 2>&1; then
    log "✓ Core network ports are listening"
else
    warn "Some core network ports may not be listening"
fi

# UPF ports
if ssh $VM2 "ss -tlnp | grep -E '(8805|2152)'" >/dev/null 2>&1; then
    log "✓ UPF ports are listening"
else
    warn "Some UPF ports may not be listening"
fi

log "5GC testbed startup completed successfully!"
log "You can now monitor the system using:"
log "  - Core logs: ssh $VM1 'journalctl -f -u open5gs-*'"
log "  - UPF logs: ssh $VM2 'journalctl -f -u open5gs-upfd'"
log "  - gNodeB logs: ssh $VM3 'tail -f gnb.log'"
log "  - UE logs: ssh $VM4 'tail -f ue.log'"

# Optional: Start monitoring
if [ "$1" = "--monitor" ]; then
    log "Starting monitoring mode..."
    echo "Press Ctrl+C to stop monitoring"
    
    # Monitor all logs simultaneously
    (
        ssh $VM1 "journalctl -f -u open5gs-amfd" &
        ssh $VM2 "journalctl -f -u open5gs-upfd" &
        ssh $VM3 "tail -f gnb.log" &
        ssh $VM4 "tail -f ue.log" &
        wait
    )
fi 